package cn.buhler.source;

import cn.buhler.jdbc.JDBCConfig;
import cn.buhler.thread.QueryThreadPoolExecutor;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.db.Db;
import cn.hutool.db.Entity;
import com.zaxxer.hikari.HikariDataSource;

import org.apache.flume.Context;
import org.apache.flume.Event;
import org.apache.flume.EventDrivenSource;
import org.apache.flume.conf.Configurable;
import org.apache.flume.conf.ConfigurationException;
import org.apache.flume.event.EventBuilder;
import org.apache.flume.source.AbstractSource;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.SQLException;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * @author shallwe_wang(xiaowei.wang@buhlergroup.com)
 * 2023年7月30日 下午2:30:39
 * 自定义数据源
 *
 */
public class RichJDBCSource extends AbstractSource implements Configurable, EventDrivenSource {
    private static final Logger LOG = LoggerFactory.getLogger(RichJDBCSource.class);
    private HikariDataSource dataSource;

    private final Map<String, DBPoller> dbPollerMap = new ConcurrentHashMap<>();

    private final QueryThreadPoolExecutor threadPoolExecutor = new QueryThreadPoolExecutor();

    /**
     * 线程正在运行的标记，用于通知线程资源的释放
     */
    private volatile boolean running;
    // agent.source.source.table.mapping = t1>t4,t2>t5,t3>t6

    private Set<String> sourceTables;

    private Map<String, String> tableMapping;



    @Override
    public void configure(Context context) {
        try {
            JDBCConfig jdbcConfig = new JDBCConfig(context);
            this.dataSource = jdbcConfig.getDataSource();
            long connectionTimeout = context.getLong("connectionTimeout", 30000L); // 默认30秒
            int maxPoolSize = context.getInteger("maxPoolSize", 5); // 默认5个连接
            this.dataSource.setConnectionTimeout(connectionTimeout);
            this.dataSource.setMaximumPoolSize(maxPoolSize);
        } catch (ConfigurationException e) {
            LOG.info(e.toString());
            throw e;
        }

        Map<String, String> tableProps = context.getSubProperties("table.");

        sourceTables = new HashSet<>();
        for (String key : tableProps.keySet()) {
            String[] parts = key.split("\\.");
            if (parts.length > 0) {
                sourceTables.add(parts[0]);
            }
        }


        String tableMappingString = context.getString("tableMapping","*");
        tableMapping =  parseMapping(tableMappingString,true);

        for (String table : sourceTables) {

            if (!tableMapping.containsKey(table)) {
                throw new ConfigurationException("Table " + table + " is not found in table.mapping configuration.");
            }
            String sql = context.getString("table."+table + ".query");
            if(StrUtil.isEmpty(sql)){
                throw new ConfigurationException(table + ".query参数为空，同步"+table+"的sql语句必须指定。");
            }
            String columnMappingString = context.getString("table."+table + ".columnMapping", "*");
            Map<String, String> columnMapping = parseMapping(columnMappingString,false);
            long interval = context.getLong("table."+table + ".poll.interval", 60L);
            DBPoller dbPoller = new DBPoller(table, columnMapping, sql, interval);
            dbPollerMap.put(table, dbPoller);
        }
    }

    /**
     * 通用的映射解析方法
     * @param mappingString 映射字符串
     * @param tableMapping  表映射还是字段映射
     * @return 映射的键值对
     */
    private Map<String, String> parseMapping(String mappingString,boolean tableMapping) {
        Map<String, String> mapping = new HashMap<>();
        if("*".equals(mappingString)){
            if(tableMapping){
                for (String key : sourceTables) {
                    mapping.put(key, key);
                }
            }else{
                return  mapping;
            }
        }
        String[] items = mappingString.split(",");
        for (String item : items) {
            if (item.contains(">")) {
                String[] parts = item.split(">");
                mapping.put(parts[0].trim(), parts[1].trim());
            } else {
                mapping.put(item.trim(), item.trim());
            }
        }

        return mapping;
    }


    @Override
    public synchronized void start() {
        super.start();
        running = true;

        for (DBPoller poller : dbPollerMap.values()) {
            threadPoolExecutor.submit(poller);
        }
        LOG.info("------------Buhler SmarT FLume RichJDBC 数据源开启成功------------");
    }

    @Override
    public synchronized void stop() {
        running = false;
        for (DBPoller poller : dbPollerMap.values()) {
            poller.stop();
        }
        threadPoolExecutor.shutdown();
        super.stop();
    }




    /**
     * 数据抽取线程
     */
    private class DBPoller implements Runnable {
        private final String sourceTable;
        private final String targetTable;
        private final String query;
        private final long interval;

        private final Map<String, String> columnMapping;

        DBPoller(String sourceTable,  Map<String, String> columnMapping, String query,  long interval) {
            this.sourceTable = sourceTable;
            this.targetTable = tableMapping.get(sourceTable);
            this.columnMapping = columnMapping;
            this.query = query;
            this.interval = interval;
            LOG.info("send targetTable: {},columnMapping :{}",sourceTable,columnMapping);
        }

        public void stop() {
            running = false;
            LOG.info("------------Buhler SmarT FLume RichJDBC 数据源准备关闭------------");
        }


        @Override
        public void run() {
            while (running) {
                List<Event> eventBatch = new ArrayList<>();
                //查询并处理数据
                try {
                    List<Entity> results = queryDatabase(query);

                    for (Entity entity : results) {
                        //默认发送所有的列
                        Entity toSend = entity;
                        //如果不是要发送所有列
                        if (!columnMapping.isEmpty()) {
                            toSend = new Entity();
                            for (Map.Entry<String, String> entry : columnMapping.entrySet()) {
                                String sourceColumn = entry.getKey();
                                String targetColumn = entry.getValue();
                                if (entity.containsKey(sourceColumn)) {
                                    toSend.put(targetColumn, entity.get(sourceColumn));
                                }else{
                                    LOG.info("从表{}中读取执行的sql{}获取的数据中并不包含字段{}",sourceTable, query, sourceColumn);
                                }
                            }
                        }
                        Event event = EventBuilder.withBody(ObjectUtil.serialize(toSend));
                        event.getHeaders().put("targetTable", targetTable);
                        eventBatch.add(event);
                    }

                } catch (SQLException e) {
                    LOG.error("从表{}中读取数据发生错误.{}", sourceTable, e);
                }
                //发生事件
                if (!eventBatch.isEmpty()) {
                    getChannelProcessor().processEventBatch(eventBatch);
                }
                try {
                    TimeUnit.SECONDS.sleep(interval);
                } catch (InterruptedException e) {
                    LOG.warn("读取数据的线程发生中断", e);
                    Thread.currentThread().interrupt();
                }
            }
        }

        private List<Entity> queryDatabase(String query) throws SQLException {
            try {
                return Db.use(dataSource).query(query);
            } catch (SQLException e) {
                throw e;
            }
        }
    }
}